import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
import yaml
import cv2
from modules.xfeat import XFeat
from evaluation.dual_softmax_matcher import DualSoftMaxMatcher
from evaluation.alike.mega_pose_est_mnn import MegaDepthPoseMNNBenchmark
from evaluation.superpoint.superpoint import SuperPoint
from third_party.ALIKE.alike import ALike

def get_best_device(verbose = False):
    device = torch.device('cpu')
    if torch.cuda.is_available():
        device = torch.device('cuda')
    elif torch.backends.mps.is_available():
        device = torch.device('mps')
    else:
        device = torch.device('cpu')
    if verbose: print (f"Fastest device found is: {device}")
    return device


if __name__ == "__main__":

    ALIKE_PATH = "/media/liyuke/share/AAA/part2/xfeat/third_party/ALIKE"
    dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    configs = {
        'alike-t': {'c1': 8, 'c2': 16, 'c3': 32, 'c4': 64, 'dim': 64, 'single_head': True, 'radius': 2,
                    'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-t.pth')},
        'alike-s': {'c1': 8, 'c2': 16, 'c3': 48, 'c4': 96, 'dim': 96, 'single_head': True, 'radius': 2,
                    'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-s.pth')},
        'alike-n': {'c1': 16, 'c2': 32, 'c3': 64, 'c4': 128, 'dim': 128, 'single_head': True, 'radius': 2,
                    'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-n.pth')},
        'alike-l': {'c1': 32, 'c2': 64, 'c3': 128, 'c4': 128, 'dim': 128, 'single_head': False, 'radius': 2,
                    'model_path': os.path.join(ALIKE_PATH, 'models', 'alike-l.pth')},
    }
    model = ALike(**configs['alike-t'],
                device=dev,
                top_k=4096,
                scores_th=0.1,
                n_limit=8000)
    matcher = DualSoftMaxMatcher()


    mega_1500 = MegaDepthPoseMNNBenchmark()
    mega_1500.benchmark(descriptor_model = model,matcher_model = matcher)